import os
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import welch
import json

# --- 图表全局设置 ---
plt.rcParams['font.sans-serif'] = ['SimSun'] # 中文使用宋体
plt.rcParams['font.serif'] = ['Times New Roman'] # 英文和数字使用新罗马
plt.rcParams['font.family'] = ['serif', 'sans-serif'] 
plt.rcParams['axes.unicode_minus'] = False

def find_displacement_data_path(root_path, region_processed_folder, flow_rate):
    """
    在指定的区域文件夹内，根据流速构建并查找 displacement_data.csv 文件的完整路径。

    参数:
        root_path (str): 数据存储的根目录。
        region_processed_folder (str): 对应测点区域的处理后文件夹全名。
        flow_rate (float): 需要分析的流速值 (单位: m/s)。

    返回:
        str: displacement_data.csv 文件的完整路径，如果未找到则返回 None。
    """
    flow_rate_str = f"{flow_rate:.1f}"
    if flow_rate == 5.0 or flow_rate == 4.1:
        flow_folder_name = f"1000f{flow_rate_str}ms"
    else:
        flow_folder_name = f"1000f{flow_rate_str}ms"

    flow_folder_path = os.path.join(root_path, region_processed_folder, flow_folder_name)

    if os.path.isdir(flow_folder_path):
        subfolders = [d for d in os.listdir(flow_folder_path) if os.path.isdir(os.path.join(flow_folder_path, d))]
        if subfolders:
            data_folder_path = os.path.join(flow_folder_path, subfolders[0])
            displacement_file_path = os.path.join(data_folder_path, 'displacement_data.csv')
            if os.path.exists(displacement_file_path):
                return displacement_file_path
            else:
                print(f"文件未找到: {displacement_file_path}")
                return None

    print(f"在 {region_processed_folder} 中未找到流速为 {flow_rate} m/s 的文件夹 (预期路径: {flow_folder_path})")
    return None

def calculate_max_amplitude_from_displacement(displacement_data):
    """
    从位移数据计算最大振幅
    
    参数:
        displacement_data (array): 位移数据数组
        
    返回:
        float: 最大振幅值
    """
    # 计算位移的均方根值作为振幅估计
    rms_amplitude = np.sqrt(np.mean(displacement_data**2))
    
    # 或者使用峰峰值的一半作为最大振幅
    peak_to_peak = np.max(displacement_data) - np.min(displacement_data)
    max_amplitude = peak_to_peak / 2
    
    return max_amplitude

def calculate_max_amplitude_from_psd(freqs, pxx, fs):
    """
    从PSD数据估算最大振幅
    
    参数:
        freqs (array): 频率数组
        pxx (array): PSD功率密度数组
        fs (float): 采样频率
        
    返回:
        float: 估算的最大振幅值
    """
    # 找到PSD的最大值对应的频率
    max_psd_idx = np.argmax(pxx)
    dominant_freq = freqs[max_psd_idx]
    max_psd_value = pxx[max_psd_idx]
    
    # 从PSD估算振幅: A = sqrt(2 * PSD * Δf)
    # 其中Δf是频率分辨率
    df = freqs[1] - freqs[0]
    estimated_amplitude = np.sqrt(2 * max_psd_value * df)
    
    return estimated_amplitude, dominant_freq

def analyze_max_amplitude_vs_flow_rate(root_path, regions, flow_rates, output_folder):
    """
    分析不同流速下的最大振幅，并生成结果图表和数据文件
    
    参数:
        root_path (str): 数据存储的根目录
        regions (list): 区域标识符列表
        flow_rates (list): 流速列表
        output_folder (str): 输出文件夹路径
    """
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    # 存储所有结果
    results = {}
    
    all_folders_in_root = os.listdir(root_path)
    
    for region_prefix in regions:
        region_results = {
            'flow_rates': [],
            'max_amplitudes_from_displacement': [],
            'max_amplitudes_from_psd': [],
            'dominant_frequencies': []
        }
        
        region_folder_name = None
        for folder in all_folders_in_root:
            if folder.startswith(region_prefix + ' ') and folder.endswith('_processed'):
                region_folder_name = folder
                break
        
        if not region_folder_name:
            print(f"警告: 在 {root_path} 中未找到以前缀 '{region_prefix}' 开头的 processed 文件夹。")
            continue
        
        for flow_rate in flow_rates:
            file_path = find_displacement_data_path(root_path, region_folder_name, flow_rate)
            
            if file_path:
                try:
                    df = pd.read_csv(file_path)
                    time_col, disp_col = df.columns[0], df.columns[1]
                    time = np.array(df[time_col])
                    displacement = np.array(df[disp_col])
                    
                    # 计算采样频率
                    fs = 1.0 / float(np.mean(np.diff(time)))
                    
                    # 方法1: 直接从位移数据计算最大振幅
                    max_amp_disp = calculate_max_amplitude_from_displacement(displacement)
                    
                    # 方法2: 从PSD估算最大振幅
                    freqs, pxx = welch(displacement, fs, nperseg=1024)
                    max_amp_psd, dom_freq = calculate_max_amplitude_from_psd(freqs, pxx, fs)
                    
                    # 存储结果
                    region_results['flow_rates'].append(flow_rate)
                    region_results['max_amplitudes_from_displacement'].append(max_amp_disp)
                    region_results['max_amplitudes_from_psd'].append(max_amp_psd)
                    region_results['dominant_frequencies'].append(dom_freq)
                    
                    print(f"区域 {region_prefix}, 流速 {flow_rate}m/s: "
                          f"位移最大振幅={max_amp_disp:.6f}, PSD估算振幅={max_amp_psd:.6f}, "
                          f"主导频率={dom_freq:.2f}Hz")
                          
                except Exception as e:
                    print(f"处理文件 {file_path} 时发生错误: {e}")
        
        results[region_prefix] = region_results
    
    # 绘制最大振幅 vs 流速图表
    plot_amplitude_vs_flow_rate(results, output_folder)
    
    # 保存结果到JSON文件
    save_results_to_json(results, output_folder)
    
    # 保存结果到CSV文件
    save_results_to_csv(results, output_folder)
    
    return results

def plot_amplitude_vs_flow_rate(results, output_folder):
    """
    绘制最大振幅随流速变化的图表
    
    参数:
        results (dict): 分析结果字典
        output_folder (str): 输出文件夹路径
    """
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 12))
    
    colors = ['black', 'magenta', 'darkviolet', 'olive', 'steelblue', 'darkcyan']
    markers = ['s', 'v', 'o', 'p', 'D', '^']
    
    for i, (region, data) in enumerate(results.items()):
        if data['flow_rates']:  # 确保有数据
            color = colors[i % len(colors)]
            marker = markers[i % len(markers)]
            
            # 绘制位移计算的最大振幅
            ax1.plot(data['flow_rates'], data['max_amplitudes_from_displacement'], 
                    color=color, marker=marker, linestyle='-', linewidth=2, markersize=8,
                    label=f'截面{region.upper()}')
            
            # 绘制PSD估算的最大振幅
            ax2.plot(data['flow_rates'], data['max_amplitudes_from_psd'], 
                    color=color, marker=marker, linestyle='-', linewidth=2, markersize=8,
                    label=f'截面{region.upper()}')
    
    # 设置第一个子图（位移计算）
    ax1.set_xlabel('流速 (m/s)', fontsize=14)
    ax1.set_ylabel('最大振幅 (从位移计算)', fontsize=14)
    ax1.set_title('最大振幅 vs 流速 (从位移数据计算)', fontsize=16)
    ax1.grid(True, alpha=0.3)
    ax1.legend(fontsize=12)
    ax1.tick_params(axis='both', which='major', labelsize=12)
    
    # 设置第二个子图（PSD估算）
    ax2.set_xlabel('流速 (m/s)', fontsize=14)
    ax2.set_ylabel('最大振幅 (从PSD估算)', fontsize=14)
    ax2.set_title('最大振幅 vs 流速 (从PSD估算)', fontsize=16)
    ax2.grid(True, alpha=0.3)
    ax2.legend(fontsize=12)
    ax2.tick_params(axis='both', which='major', labelsize=12)
    
    plt.tight_layout()
    
    # 保存图表
    output_path = os.path.join(output_folder, 'Max_Amplitude_vs_Flow_Rate.png')
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"最大振幅vs流速图表已保存至: {output_path}")

def save_results_to_json(results, output_folder):
    """
    将结果保存为JSON文件
    
    参数:
        results (dict): 分析结果字典
        output_folder (str): 输出文件夹路径
    """
    # 转换numpy数组为Python列表以便JSON序列化
    json_results = {}
    for region, data in results.items():
        json_results[region] = {
            'flow_rates': [float(x) for x in data['flow_rates']],
            'max_amplitudes_from_displacement': [float(x) for x in data['max_amplitudes_from_displacement']],
            'max_amplitudes_from_psd': [float(x) for x in data['max_amplitudes_from_psd']],
            'dominant_frequencies': [float(x) for x in data['dominant_frequencies']]
        }
    
    output_path = os.path.join(output_folder, 'max_amplitude_results.json')
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(json_results, f, indent=2, ensure_ascii=False)
    print(f"JSON结果文件已保存至: {output_path}")

def save_results_to_csv(results, output_folder):
    """
    将结果保存为CSV文件
    
    参数:
        results (dict): 分析结果字典
        output_folder (str): 输出文件夹路径
    """
    all_data = []
    
    for region, data in results.items():
        for i in range(len(data['flow_rates'])):
            row = {
                'region': region,
                'flow_rate': data['flow_rates'][i],
                'max_amplitude_from_displacement': data['max_amplitudes_from_displacement'][i],
                'max_amplitude_from_psd': data['max_amplitudes_from_psd'][i],
                'dominant_frequency': data['dominant_frequencies'][i]
            }
            all_data.append(row)
    
    if all_data:
        df = pd.DataFrame(all_data)
        output_path = os.path.join(output_folder, 'max_amplitude_results.csv')
        df.to_csv(output_path, index=False, encoding='utf-8-sig')
        print(f"CSV结果文件已保存至: {output_path}")

if __name__ == '__main__':
    # --- 用户配置区域 ---
    ROOT_DATA_PATH = r'J:\lqb\huanxingranliao\25CJ02'
    OUTPUT_FOLDER = 'max_amplitude_analysis_results'
    regions_to_analyze = ['A', 'B', 'C']
    flow_rates_to_analyze = [0.1,0.3,0.6,0.9, 1.2,1.5,1.8, 2.0,2.2,2.4,2.6,2.8, 3.0,3.2,3.5,3.8, 4.1,4.4,4.7, 5.0]

    # --- 执行最大振幅分析 ---
    print("开始分析最大振幅数据...")
    results = analyze_max_amplitude_vs_flow_rate(
        ROOT_DATA_PATH, regions_to_analyze, flow_rates_to_analyze, OUTPUT_FOLDER
    )
    print("分析完成！")
